#include <iostream>
#include <vector>
#include <unordered_map>
#include <queue>
#include <cassert>

using namespace std;

// 347.前K个高频元素
// https://leetcode.cn/problems/top-k-frequent-elements/description/
// 给定一个非空数组，返回前k个出现频率最高的元素。
// 示例 1:
// 输入: nums = [1,1,1,2,2,3], k = 2
// 输出: [1,2]

class Solution {
public:
    vector<int> topKFrequent(vector<int>& nums, int k) {
        assert(k > 0);

        // 统计每个元素出现的频率
        unordered_map<int, int> freq;
        for(int i = 0; i < nums.size(); i++)
            freq[nums[i]]++;

        assert(k <= freq.size());

        // 扫描 freq，维护当前出现频率最高的k个元素
        // 在优先队列中，按照频率排序，所以数据对象是（频率，元素）的形式。
        priority_queue<pair<int, int>, vector<pair<int, int>>, greater<pair<int,int>> > pq; // 底层为最小堆的优先队列
        for(auto iter = freq.begin(); iter != freq.end(); iter++) {
            // 优先队列已经慢了
            if(pq.size() == k) {
                if(iter->second > pq.top().first) {
                    // 取代优先队列中频率最小的元素
                    pq.pop();
                    pq.push(make_pair(iter->second, iter->first));
                }
            }
            else
                pq.push(make_pair(iter->second,iter->first));
        }

        vector<int> res;
        while(!pq.empty()) {
            res.push_back(pq.top().second);
            pq.pop();
        }

        return res;
    }
};

int main() {
    int nums[] = {1, 1, 1, 2, 2, 3};
    vector<int> vec(nums, nums + sizeof(nums)/sizeof(int));
    int k = 2;

    vector<int> res = Solution().topKFrequent(vec, 2);
    for( int i = 0 ; i < res.size() ; i ++ )
        cout<<res[i]<<" ";
    cout<<endl;
    return 0;
}
